from numpy import *
def loadDataSet(fileName):  
    """读取数据

    Args:
        fileName (string): 文件名称

    Returns:
        list: 购物篮
    """  
    dataMat = []                
    fr = open(fileName)
    for line in fr.readlines():
        curLine = line.strip().split(',')
        dataMat.append(curLine)
    return dataMat


# 一、频繁项集产生

def createC1(dataSet):
    """产生候选1-项集

    Args:
        dataSet (list): 购物篮

    Returns:
        list(frozenset): 1-项集
    """
    C1 = [] #用于记录候选1-项集
    for transaction in dataSet:
        for item in transaction:
            if [item] not in C1:
                C1.append([item])
    C1.sort()
    return list(map(frozenset, C1)) #以frozenset的形式存储每一个项集


def scanD(D, Ck, minSupport):
    """计算所有候选k-项集的支持度，找到所有频繁k-项集

    Args:
        D (list(set)): 购物篮数据集
        Ck (list(frozenset)): 候选k-项集
        minSupport (float): 支持度阈值

    Returns:
        list, dictionary: 达到支持度阈值的所有k-项集及其支持度
    """
    ssCnt = {}  #k-项集对应的支持度计数(dictionary)
    numItems = 0.0  #记录购物篮的事务总数(float),这里不能直接用len()求D的长度
    for tid in D:
        for can in Ck:
            if can.issubset(tid):   #该k-项集(can)是否包含在这个事务(tid)中
                if can not in ssCnt:
                    ssCnt[can] = 1
                else: 
                    ssCnt[can] += 1
        numItems += 1
    retList = []    #记录达到支持度阈值的k-项集
    supportData = {}    #记录达到支持度阈值的k-项集的支持度
    for key in ssCnt:
        support = ssCnt[key]/numItems   #支持度
        #print(key, support)
        if support >= minSupport:
            retList.insert(0,key)
            supportData[key] = support  #存储达到支持度阈值候选项集的支持度
    return retList, supportData

def aprioriGen(Fk, k): #creates Ck
    """由k-项集生成候选(k+1)-项集:F_k-1 * F_k-1

    Args:
        Fk (list(set)): 频繁k-项集
        k (int)): 当前频繁项集的长度+1

    Returns:
        list(set): 该k-项集对应的候选(k+1)-项集集合
    """
    retList = []
    lenFk = len(Fk)
    for i in range(lenFk):
        for j in range(i+1, lenFk): 
            L1 = list(Fk[i]); L2 = list(Fk[j])
            L1.sort(); L2.sort()    #排序操作
            if L1[:k-2]==L2[:k-2]:
                retList.append(Fk[i] | Fk[j]) #取并集，加入候选项集
    return retList

def apriori(dataSet, minSupport):
    """生成频繁项集

    Args:
        dataSet (list): 购物篮数据集
        minSupport (float): 支持度阈值

    Returns:
        list(set), dictionary: 所有的频繁项集, 频繁项集的支持度 
    """
    C1 = createC1(dataSet)
    D = list(map(set, dataSet))   #将每一个事务转换为set
    F1, supportData = scanD(D, C1, minSupport)
    F = [F1]    #频繁1-项集整体作为L的一项，L用来记录所有的频繁项集，list的第一维表示频繁项集的k
    k = 2
    """
    下面需要注意：
    k=2时，L[k-2]表示所有的频繁1-项集，也就是L1；
    k=2意味着接下来要从频繁1-项集生成候选2-项集
    """
    while (len(F[k-2]) > 0):    #当频繁(k-1)-项集不为空时，需要挖掘这些频繁项集的超集
        Ck = aprioriGen(F[k-2], k)  #由频繁(k-1)-项集扩展到候选k-项集（所有的）：
        Lk, supK = scanD(D, Ck, minSupport) #Lk表示所有频繁k-项集的集合（type:list(set)）
        supportData.update(supK)
        F.append(Lk)    #将所有频繁k-项集的集合加入到L的末尾
        k += 1
    return F[:-1], supportData  #L的最后一个频繁项集没有被return，因为是空的，用于跳出while循环


# 二、规则产生

def generateRules(F, supportData, minConf):  #supportData is a dict coming from scanD
    """规则产生

    Args:
        F (list(frozenset)): 所有的频繁项集
        supportData (dictionary): 频繁项集的支持度
        minConf (float, optional): 置信度阈值

    Returns:
        list(frozenset): 关联规则
    """
    bigRuleList = []
    
    #只选取频繁项集长度大于等于2的,这样的需要产生规则
    for i in range(1, len(F)):
        for freqSet in F[i]:
            H1 = [frozenset([item]) for item in freqSet] #首先生成后件（包含该频繁项集的所有商品）
            #使用置信度度量对关联规则进行剪枝
            #对于频繁2-项集，只需要进行一次前后件的拆分来组成规则；而对于频繁k-项集(k>2)时，需要做剪枝操作
            if i > 1:   
                #从频繁3-项集开始，才需要进行剪枝的操作（首先要将后件的长度变长）
                rulesFromConseq(freqSet, H1, supportData, bigRuleList, minConf) 
            else:
                calcConf(freqSet, H1, supportData, bigRuleList, minConf)    #运行该函数更新了bigRuleList
    return bigRuleList         

def calcConf(freqSet, H, supportData, brl, minConf):
    """计算k-频繁项集的置信度

    Args:
        freqSet (frozenset): 一个频繁k-项集
        H (list(frozenset)): 该频繁项集下的候选规则后件
        supportData (dictionary): 候选项集的支持度
        brl (list): 存储关联规则
        minConf (float, optional): 置信度阈值

    Returns:
        list: 达到置信度阈值的该频繁k-项集下的后件，
                例如频繁项集{a,b,c}，其中b,c做长度为1的规则后件时满足minConf，则return{b,c}
    """
    #prunedH = []    #记录该频繁k-项集满足置信度阈值的m-后件集合
    for conseq in H:
        conf = supportData[freqSet]/supportData[freqSet-conseq] #计算置信度
        if conf >= minConf: 
            brl.append((freqSet-conseq, conseq, conf))  #存储关联规则
            #prunedH.append(conseq)
    #return prunedH

def rulesFromConseq(freqSet, H, supportData, brl, minConf):
    """递归生成后件，计算每一个规则的置信度，并做剪枝操作

    Args:
        freqSet (frozenset): 一个频繁k-项集，即在从该结点出发来剪枝
        H (list(frozenset)): 该频繁项集下的候选规则后件
        supportData (dictionary): 候选项集的支持度
        brl (list)): 存储关联规则
        minConf (float, optional): 置信度阈值
    """
    m = len(H[0])   #后件的长度
    if (len(freqSet) > m):  
        Hmp1 = calcConf(freqSet, H, supportData, brl, minConf)  
        Hmp1 = aprioriGen(H, m+1)   #用aprioriGen函数将后件的长度扩展，即寻找该频繁k-项集对应的频繁项集
        if (len(Hmp1) > 1):    
            rulesFromConseq(freqSet, Hmp1, supportData, brl, minConf)
     
def pntRules(ruleList):
    """输出关联规则

    Args:
        ruleList (list): 关联规则
    """
    
    for ruleTup in ruleList:
        for item in ruleTup[0]:
            print("?")
            print(item, end=" ")
        print(" --------> ", end="")
        for item in ruleTup[1]:
            print(item, end=" ")
        print("     confidence: %f" % ruleTup[2])

# 主函数

def main():
    minSupport = 0.04
    minConf = 0.06
    dataSet = loadDataSet('Groceries.txt')
    F, supportData = apriori(dataSet, minSupport)
    ruleList = generateRules(F, supportData, minConf)
    pntRules(ruleList)  #打印关联规则
    
    
if __name__ == '__main__':
    main()
